#include "ompl/contrib/hplanner/StateTransformer.h"

namespace ompl
{
namespace base
{


StateToVectorTransformer::StateToVectorTransformer(const SpaceInformationPtr& si) : space_(si->getStateSpace())
{
  dim_ = 0;
  computeDim(space_, dim_);

  if(dim_ > 0)
  {
    cache_ = new float[dim_];
    upper_bound_ = new float[dim_];
    lower_bound_ = new float[dim_];
    weights_ = new float[dim_];
    types_ = new bool[dim_];

    for(unsigned int i = 0; i < dim_; ++i)
    {
      cache_[i] = 0;
      upper_bound_[i] = 0;
      lower_bound_[i] = 0;
      weights_[i] = 0;
      types_[i] = true;

    }

    unsigned int offset = 0;
    float weight = 1;
    getData(space_, offset);
    offset = 0;
    getWeights(space_, offset, weight);
  }
  else
  {
    cache_ = NULL;
    upper_bound_ = NULL;
    lower_bound_ = NULL;
    weights_ = NULL;
  }
}


StateToVectorTransformer::~StateToVectorTransformer()
{
  delete [] cache_;
  delete [] upper_bound_;
  delete [] lower_bound_;
  delete [] weights_;
  delete [] types_;
}

void StateToVectorTransformer::transform(const State* state, StateSpacePtr space, unsigned int& offset) const
{
  if(space->isCompound())
  {
    const std::vector<StateSpacePtr>& subspaces = static_cast<const CompoundStateSpace*>(space.get())->getSubSpaces();

    const CompoundState* cstate = static_cast<const CompoundState*>(state);
    for(unsigned int i = 0; i < subspaces.size(); ++i)
    {
      transform((*cstate)[i], subspaces[i], offset);
    }
  }
  else
  {
    switch(space->getType())
    {
      case STATE_SPACE_SE3:
      {
        const typename SE3StateSpace::StateType* s = static_cast<const typename SE3StateSpace::StateType *>(state);
        cache_[0 + offset] = s->getX();
        cache_[1 + offset] = s->getY();
        cache_[2 + offset] = s->getZ();

        double a, b, c;
        quaternionToEulerAngles(s->rotation().x, s->rotation().y, s->rotation().z, s->rotation().w, a, b, c);
        cache_[3 + offset] = a;
        cache_[4 + offset] = b;
        cache_[5 + offset] = c;
        offset += 6;
      }
      break;
      case STATE_SPACE_SE2:
      {
        const typename SE2StateSpace::StateType* s = static_cast<const typename SE2StateSpace::StateType *>(state);
        cache_[0 + offset] = s->getX();
        cache_[1 + offset] = s->getY();
        cache_[2 + offset] = s->getYaw();
        offset += 3;
      }
      break;
      case STATE_SPACE_SO3:
      {
        const typename SO3StateSpace::StateType* s = static_cast<const typename SO3StateSpace::StateType *>(state);
        double a, b, c;
        quaternionToEulerAngles(s->x, s->y, s->z, s->w, a, b, c);
        cache_[0 + offset] = a;
        cache_[1 + offset] = b;
        cache_[2 + offset] = c;
        offset += 3;
      }
      break;
      case STATE_SPACE_SO2:
      {
        const typename SO2StateSpace::StateType* s = static_cast<const typename SO2StateSpace::StateType *>(state);
        cache_[0 + offset] = s->value;
        offset += 1;
      }
      break;
      case STATE_SPACE_REAL_VECTOR:
      {
        const typename RealVectorStateSpace::StateType* s = static_cast<const typename RealVectorStateSpace::StateType *>(state);

        for(unsigned int i = 0; i < space->getDimension(); ++i)
        {
          cache_[i + offset] = s->values[i];
        }
        offset += space->getDimension();
      }
      break;
      default:
        std::cout << "space type not supported" << std::endl;
    }
  }
}


void StateToVectorTransformer::computeDim(StateSpacePtr space, unsigned int& dim)
{
  if(space->isCompound())
  {
    const std::vector<StateSpacePtr>& subspaces = static_cast<const CompoundStateSpace*>(space.get())->getSubSpaces();

    for(unsigned int i = 0; i < subspaces.size(); ++i)
    {
      computeDim(subspaces[i], dim);
    }
  }
  else
  {
    switch(space->getType())
    {
      case STATE_SPACE_SE3:
        dim += 6;
        break;
      case STATE_SPACE_SE2:
        dim += 3;
        break;
      case STATE_SPACE_SO3:
        dim += 3;
        break;
      case STATE_SPACE_SO2:
        dim += 1;
        break;
      case STATE_SPACE_REAL_VECTOR:
        dim += space->getDimension();
        break;
      default:
        std::cout << "space type not supported" << std::endl;
    }
  }
}


void StateToVectorTransformer::getData(StateSpacePtr space, unsigned int& offset)
{
  if(space->isCompound())
  {
    const std::vector<StateSpacePtr>& subspaces = static_cast<const CompoundStateSpace*>(space.get())->getSubSpaces();

    for(unsigned int i = 0; i < subspaces.size(); ++i)
    {
      getData(subspaces[i], offset);
    }
  }
  else
  {
    switch(space->getType())
    {
      case STATE_SPACE_SE3:
      {
        RealVectorBounds bounds = static_cast<const SE3StateSpace*>(space.get())->getBounds();
        upper_bound_[0 + offset] = bounds.high[0];
        upper_bound_[1 + offset] = bounds.high[1];
        upper_bound_[2 + offset] = bounds.high[2];

        lower_bound_[0 + offset] = bounds.low[0];
        lower_bound_[1 + offset] = bounds.low[1];
        lower_bound_[2 + offset] = bounds.low[2];

        upper_bound_[3 + offset] = 3.1415926;
        upper_bound_[4 + offset] = 3.1415926;
        upper_bound_[5 + offset] = 3.1415926;

        lower_bound_[3 + offset] = -3.1415926;
        lower_bound_[4 + offset] = -3.1415926;
        lower_bound_[5 + offset] = -3.1415926;

        types_[0 + offset] = true;
        types_[1 + offset] = true;
        types_[2 + offset] = true;
        types_[3 + offset] = false;
        types_[4 + offset] = false;
        types_[5 + offset] = false;
        offset += 6;
      }
      break;
      case STATE_SPACE_SE2:
      {
        RealVectorBounds bounds = static_cast<const SE2StateSpace*>(space.get())->getBounds();
        upper_bound_[0 + offset] = bounds.high[0];
        upper_bound_[1 + offset] = bounds.high[1];

        lower_bound_[0 + offset] = bounds.low[0];
        lower_bound_[1 + offset] = bounds.low[1];

        upper_bound_[2 + offset] = 3.1415926;
        lower_bound_[2 + offset] = -3.1415926;

        types_[0 + offset] = true;
        types_[1 + offset] = true;
        types_[2 + offset] = false;

        offset += 3;
      }
      break;
      case STATE_SPACE_SO3:
      {
        upper_bound_[0 + offset] = 3.1415926;
        upper_bound_[1 + offset] = 3.1415926;
        upper_bound_[2 + offset] = 3.1415926;

        lower_bound_[0 + offset] = -3.1415926;
        lower_bound_[1 + offset] = -3.1415926;
        lower_bound_[2 + offset] = -3.1415926;

        types_[0 + offset] = false;
        types_[1 + offset] = false;
        types_[2 + offset] = false;

        offset += 3;
      }
      break;
      case STATE_SPACE_SO2:
      {
        upper_bound_[0 + offset] = 3.1415926;
        lower_bound_[0 + offset] = -3.1415926;

        types_[0 + offset] = false;

        offset += 1;
      }
      break;
      case STATE_SPACE_REAL_VECTOR:
      {
        RealVectorBounds bounds = static_cast<const RealVectorStateSpace*>(space.get())->getBounds();
        for(unsigned int i = 0; i < bounds.low.size(); ++i)
        {
          upper_bound_[i + offset] = bounds.high[i];
          lower_bound_[i + offset] = bounds.low[i];
          types_[i + offset] = true;
        }
        offset += bounds.low.size();
      }
      break;
      default:
        std::cout << "space type not supported" << std::endl;
    }
  }
}

void StateToVectorTransformer::getWeights(StateSpacePtr space, unsigned int& offset, float weight)
{
  if(space->isCompound())
  {
    const CompoundStateSpace* s = static_cast<const CompoundStateSpace*>(space.get());
    const std::vector<StateSpacePtr>& subspaces = s->getSubSpaces();

    for(unsigned int i = 0; i < subspaces.size(); ++i)
    {
      getWeights(subspaces[i], offset, weight * s->getSubSpaceWeight(i));
    }
  }
  else
  {
    switch(space->getType())
    {
      case STATE_SPACE_SE3:
      {
        RealVectorBounds bounds = static_cast<const SE3StateSpace*>(space.get())->getBounds();
        weights_[0 + offset] = weight;
        weights_[1 + offset] = weight;
        weights_[2 + offset] = weight;

        weights_[3 + offset] = weight;
        weights_[4 + offset] = weight;
        weights_[5 + offset] = weight;

        offset += 6;
      }
      break;
      case STATE_SPACE_SE2:
      {
        RealVectorBounds bounds = static_cast<const SE2StateSpace*>(space.get())->getBounds();
        weights_[0 + offset] = weight;
        weights_[1 + offset] = weight;
        weights_[2 + offset] = weight;
        offset += 3;
      }
      break;
      case STATE_SPACE_SO3:
      {
        weights_[0 + offset] = weight;
        weights_[1 + offset] = weight;
        weights_[2 + offset] = weight;
        offset += 3;
      }
      break;
      case STATE_SPACE_SO2:
      {
        weights_[0 + offset] = weight;
        offset += 1;
      }
      break;
      case STATE_SPACE_REAL_VECTOR:
      {
        RealVectorBounds bounds = static_cast<const RealVectorStateSpace*>(space.get())->getBounds();
        for(unsigned int i = 0; i < bounds.low.size(); ++i)
        {
          weights_[i + offset] = weight;
        }
        offset += bounds.low.size();
      }
      break;
      default:
        std::cout << "space type not supported" << std::endl;
    }
  }
}

}
}
